/* Copyright 2016 Michael Sladoje and Mike Schälchli. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

package ch.zhaw.facerecognitionlibrary.Recognition;

import android.content.Context;
import android.graphics.Bitmap;

import org.opencv.android.Utils;
import org.opencv.core.Mat;
import org.opencv.core.Size;
import org.opencv.imgproc.Imgproc;
import org.opencv.utils.Converters;
import org.tensorflow.contrib.android.TensorFlowInferenceInterface;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

import ch.zhaw.facerecognitionlibrary.Helpers.FileHelper;
import ch.zhaw.facerecognitionlibrary.Helpers.PreferencesHelper;

/***************************************************************************************
 *    Title: TensorFlowAndroidDemo
 *    Author: miyosuda
 *    Date: 23.04.2016
 *    Code version: -
 *    Availability: https://github.com
 *
 ***************************************************************************************/

/**
 * TensorFlow with SVM or KNN
 */
public class TensorFlow implements Recognition {
    private String inputLayer; // 输入层
    private String outputLayer; // 输出层

    private int inputSize; // 输入大小
    private int channels; // 输入通道数量
    private int imageMean; // 图像的平均值
    private int imageStd; // 图片标准
    private int outputSize; // 输出大小

    private Recognition rec;

    private TensorFlowInferenceInterface inferenceInterface;

    private boolean logStats = false;

    public TensorFlow(Context context, int method) {
        String dataPath = FileHelper.TENSORFLOW_PATH;
        PreferencesHelper preferencesHelper = new PreferencesHelper(context);
        inputSize = preferencesHelper.getTensorFlowInputSize(); // 默认 160
        channels = preferencesHelper.getTensorFlowInputChannels(); // 默认 3
        imageMean = preferencesHelper.getTensorFlowImageMean(); // 默认 128
        imageStd = preferencesHelper.getTensorFlowImageStd(); // 默认 128
        outputSize = preferencesHelper.getTensorFlowOutputSize(); // 默认 128
        inputLayer = preferencesHelper.getTensorFlowInputLayer(); // 默认 input
        outputLayer = preferencesHelper.getTensorFlowOutputLayer(); // 默认 embeddings

        String modelFile = preferencesHelper.getTensorFlowModelFile(); // 默认 optimized_facenet.pb
        Boolean classificationMethod = preferencesHelper.getClassificationMethodTFCaffe(); // TensorFlow的分类方法

        // Use internal assets file as fallback, if no model file is provided
        File file = new File(dataPath + modelFile);
        if(file.exists()){
            inferenceInterface = new TensorFlowInferenceInterface(context.getAssets(), dataPath + modelFile); // 加载模型
        } else {
            inferenceInterface = new TensorFlowInferenceInterface(context.getAssets(), modelFile); // 加载模型
        }

        if(classificationMethod){
            rec = new SupportVectorMachine(context, method); // 默认
        }
        else {
            rec = new KNearestNeighbor(context, method);
        }
    }

    public TensorFlow(Context context, int inputSize, int outputSize, String inputLayer, String outputLayer, String modelFile){
        this.inputSize = inputSize;
        this.outputSize = outputSize;
        this.inputLayer = inputLayer;
        this.outputLayer = outputLayer;

        inferenceInterface = new TensorFlowInferenceInterface(context.getAssets(), modelFile); // 加载模型
    }

    /**
     * 训练
     * @return
     */
    @Override
    public boolean train() {
        return rec.train();
    }

    /**
     * 识别
     * @param img
     * @param expectedLabel 预期标签
     * @return
     */
    @Override
    public String recognize(Mat img, String expectedLabel) {
        return rec.recognize(getFeatureVector(img), expectedLabel);
    }

    /**
     * 保存到文件
     */
    @Override
    public void saveToFile() {

    }

    /**
     * 加载文件
     */
    @Override
    public void loadFromFile() {

    }

    /**
     * 保存测试数据
     */
    @Override
    public void saveTestData() {
        rec.saveTestData();
    }

    /**
     * 添加图片
     * @param img 图片
     * @param label 标签
     * @param featuresAlreadyExtracted  特征向量是否已被提取
     */
    @Override
    public void addImage(Mat img, String label, boolean featuresAlreadyExtracted) {
        if (featuresAlreadyExtracted){
            rec.addImage(img, label, true);
        } else {
            rec.addImage(getFeatureVector(img), label, true);
        }
    }

    /**
     * 获取特征向量的图片矩阵
     * @param img
     * @return
     */
    @Override
    public Mat getFeatureVector(Mat img){
        Imgproc.resize(img, img, new Size(inputSize, inputSize));

        inferenceInterface.feed(inputLayer, getPixels(img), 1, inputSize, inputSize, channels); // 将输入数据复制到TensorFlow中
        inferenceInterface.run(new String[]{outputLayer}, logStats); // 运行推理
        float[] outputs = new float[outputSize];
        inferenceInterface.fetch(outputLayer, outputs); // 将输出张量复制回输出数组

        List<Float> fVector = new ArrayList<>();
        for(float o : outputs){
            fVector.add(o);
        }

        return Converters.vector_float_to_Mat(fVector); // 将向量转为图片矩阵
    }

    /**
     * 获取图片像素
     * @param img
     * @return
     */
    private float[] getPixels(Mat img){
        Bitmap bmp = Bitmap.createBitmap(inputSize, inputSize, Bitmap.Config.ARGB_8888);
        Utils.matToBitmap(img, bmp);
        int[] intValues = new int[inputSize * inputSize];
        bmp.getPixels(intValues, 0, inputSize, 0, 0, inputSize, inputSize);

        float[] floatValues = new float[inputSize * inputSize * channels];
        for (int i = 0; i < intValues.length; ++i) {
            final int val = intValues[i];
            floatValues[i * 3 + 0] = (((float)((val >> 16) & 0xFF)) - imageMean) / imageStd; // ARGB_8888 中的 R
            floatValues[i * 3 + 1] = (((float)((val >> 8) & 0xFF)) - imageMean) / imageStd; // ARGB_8888 中的 G
            floatValues[i * 3 + 2] = (((float)(val & 0xFF)) - imageMean) / imageStd; // ARGB_8888 中的 B
        }

        return floatValues;
    }
}
